#include "scene.h"

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <limits>
#include "material.h"
#include "primitive.h"
#include "ray.h"

namespace {

const uint32_t kRussianRouletteDepth = 5;
const uint32_t kMaxDepth = 32;

const raytrace::Vector3 kBackColor(0, 0, 0);

} // end namespace

namespace raytrace {

Scene::Scene()
  : mode_(MODE_WHITTED),
    // TODO: data drive
    num_light_samples_(1) {
}

Scene::~Scene() {
  Clear();
}

void Scene::Clear() {
  for (PrimitiveList::iterator it = primitives_.begin();
       it != primitives_.end(); ++it)
    delete *it;
  primitives_.clear();
  lights_.clear();
}

void Scene::DumpToLog() const {
  printf("+++ Primitives +++\n");
  for (PrimitiveList::const_iterator it = primitives_.begin();
       it != primitives_.end(); ++it) {
    (*it)->DumpToLog();
  }

  printf("+++ Lights +++\n");
  for (PrimitiveList::const_iterator it = lights_.begin();
       it != lights_.end(); ++it) {
    (*it)->DumpToLog();
  }
}

void Scene::AddPrimitive(Primitive* prim) {
  prim->id_ = primitives_.size();
  primitives_.push_back(prim);
  if (prim->material()->emittance_color_.length_sqr() > 0) {
    lights_.push_back(prim);
  }
}

Vector3 Scene::Trace(const Ray& r) const {
  switch (mode_) {
    case MODE_WHITTED:
    case MODE_DEBUG_NORMAL:
      return TraceRay(r, 0);

    case MODE_REVERSE_PATH:
    case MODE_DEBUG_PATH_COUNT:
      return TracePath(r, 0);

    default:
      assert(false && "Unexpected mode.");
      break;
  };

  return Vector3(1, 0, 1);
}

// TODO: find out why the light dot is < 0 between floor and light
// TODO: add support for spec/refract
Vector3 Scene::TraceRay(const Ray& r, uint32_t depth) const {
  // Find the nearest primitive
  float min_d = std::numeric_limits<float>::infinity();
  PrimitiveList::const_iterator min_prim_it = Intersect(r, min_d);

  // Calculate the color
  if (min_prim_it == primitives_.end())
    return kBackColor;

  Vector3 int_pt = Vector3::add(r.o_, Vector3::scale(r.d_, min_d));
    
  const Primitive* min_prim = *min_prim_it;
  const Vector3 N = min_prim->NormalAt(int_pt);

  if (mode_ == MODE_DEBUG_NORMAL)
    return Vector3::scale(Vector3::add(N, Vector3(1,1,1)), 0.5f);

  //printf("    intersect [%lu] @ %.3f %.3f %.3f\n",
  //       min_prim->id_, int_pt.x_, int_pt.y_, int_pt.z_);
  
  Vector3 color = min_prim->material()->emittance_color_;

  //printf("    checking lights\n");
  for (PrimitiveList::const_iterator light_it = lights_.begin();
       light_it != lights_.end(); ++light_it) {
    const Primitive* light = *light_it;

    Vector3 L = Vector3::sub(light->pos_, int_pt);
    L.normalize();

    // Check if we there are any primitives between intersection and light.
    float visible_factor = 1.0f;
    /*
    const Vector3 shadow_ray_o =
          Vector3::add(int_pt,
                       Vector3::scale(L,
                                      std::numeric_limits<float>::epsilon()));
   const Ray shadow_ray(shadow_ray_o, L);
   float shadow_min_d = std::numeric_limits<float>::infinity();
   PrimitiveList::const_iterator shadow_prim_it = Intersect(shadow_ray,
                                                            shadow_min_d);

    if (shadow_prim_it != min_prim_it)
      visible_factor = 0.0f;
*/
    // Calculate the diffuse light contribution.
    const float dot = Vector3::dot(N, L);
    //printf("    dot: %.3f %.3f %.3f . %.3f %.3f %.3f -> %.3f\n",
    //       N.x_, N.y_, N.z_, L.x_, L.y_, L.z_, dot);
    if (dot > 0) {
      const Vector3& r = min_prim->material()->reflectance_color_;
      const Vector3& e = light->material()->emittance_color_;
      //printf("    r: %.3f %.3f %.3f | e: %.3f %.3f %.3f\n",
      //       r.x_, r.y_, r.z_, e.x_, e.y_, e.z_);
      Vector3 cont = Vector3::scale(dot * visible_factor,
                        Vector3::mul(r, e));
      color = Vector3::add(color, Vector3::scale(cont, 1.0f / lights_.size()));
    }
  }

//  printf("color: %.3f %.3f %.3f\n", color.x_, color.y_, color.z_);

  return color;
}

Vector3 Scene::TracePath(const Ray& r, uint32_t depth) const {
  // Find the nearest primitive
  float min_d = std::numeric_limits<float>::infinity();
  PrimitiveList::const_iterator min_prim_it = Intersect(r, min_d);
 
  if (min_prim_it == primitives_.end()) {
    return kBackColor;
  }

  // Start a new ray
  const Vector3 int_pt = Vector3::add(r.o_, Vector3::scale(r.d_, min_d));
  const Vector3 normal = (*min_prim_it)->NormalAt(int_pt);

  const Material* const m = (*min_prim_it)->material();

  if (depth > kMaxDepth)
    return Vector3(1.0f, 0.0f, 1.0f); //m->emittance_color_;

  // Russian roulette
  Vector3 reflectance_color = m->reflectance_color_;
  if (depth > kRussianRouletteDepth) {
    // max component of reflectance
    const float p = reflectance_color.x_ > reflectance_color.y_ &&
                    reflectance_color.x_ > reflectance_color.z_ ?
                      reflectance_color.x_ :
                      reflectance_color.y_ > reflectance_color.z_ ?
                        reflectance_color.y_ : reflectance_color.z_;
    if (p > 0 && drand48() < p) {
      reflectance_color = Vector3::scale(reflectance_color, 1.0f / p);
    } else {
      if (mode_ == MODE_DEBUG_PATH_COUNT)
        return Vector3(0, 1, 0);
      else
        return m->emittance_color_;
    }
  }

  // Sample light.
  Vector3 incoming_light_color;
  /*
  for (uint32_t light_sample = 0;
       light_sample < num_light_samples_; 
       ++light_sample) {
    uint32_t light_index = (uint32_t) (drand48() * lights_.size());
    Primitive* light = lights_[light_index];
    const Vector3 light_pt = light->GetRandomPoint();

    Vector3 light_ray_dir = Vector3::sub(light_pt, int_pt);
    light_ray_dir.normalize();

    Vector3 light_ray_origin =
            Vector3::add(int_pt,
                         Vector3::scale(light_ray_dir,
                                        std::numeric_limits<float>::epsilon()));
    Ray light_ray(light_ray_origin, light_ray_dir);
    float light_d;
    PrimitiveList::const_iterator light_it = Intersect(light_ray, light_d);
    if (light == *light_it) {
      // TODO: get probability here from material
      const float cos_omega = Vector3::dot(light_ray_dir, normal);
      incoming_light_color =
          Vector3::add(incoming_light_color,
                       Vector3::scale(light->material()->emittance_color_,
                                      cos_omega));
    }
  }
  incoming_light_color = Vector3::scale(incoming_light_color,
                                        1.0f / float(num_light_samples_));
  */
  // Generate reflected ray.
  std::vector<Vector3> outgoings;
  std::vector<float> ipdfs;

  m->GetOutDirection(r.d_, normal, depth, ipdfs, outgoings);
  if (ipdfs.size() != outgoings.size()) {
    printf("ERROR: Size mismatch: %u %u for material %s\n",
           ipdfs.size(), outgoings.size(), m->type());
    assert(false);
  }
  
  Vector3 incoming_color;

  // TODO: MODE_DEBUG_PATH_COUNT

  for (uint32_t i = 0; i < outgoings.size(); ++i) {
      outgoings[i].normalize();

      // Send out the reflection ray
      const Vector3 reflect_o = Vector3::add(int_pt,
                          Vector3::scale(outgoings[i],
                                        std::numeric_limits<float>::epsilon()));
      const Ray reflect_ray(reflect_o, outgoings[i]);

      const Vector3 reflected = TracePath(reflect_ray, depth + 1);
      const float f = ipdfs[i] / (2 * M_PI);
      const float cont = f / outgoings.size();

      incoming_color = Vector3::add(incoming_color,
                                    Vector3::scale(reflected, cont));
  }

  Vector3 path_color = m->emittance_color_;
  path_color = Vector3::add(path_color, incoming_light_color);
  path_color = Vector3::add(path_color, Vector3::mul(reflectance_color,
                                                     incoming_color));
/*
  if (path_color.x_ < 0.0f || path_color.x_ > 1.0f ||
      path_color.y_ < 0.0f || path_color.y_ > 1.0f ||
      path_color.z_ < 0.0f || path_color.z_ > 1.0f) {
    printf("WARNING: trace returning overflow @ %lu, %.3f %.3f %.3f\n", depth,
           path_color.x_, path_color.y_, path_color.z_);
    printf("    m: %s\n", m->type());
    printf("    e: %.3f %.3f %.3f\n",
           m->emittance_color_.x_,
           m->emittance_color_.y_,
           m->emittance_color_.z_);
    printf("    r: %.3f %.3f %.3f\n",
           reflectance_color.x_,
           reflectance_color.y_,
           reflectance_color.z_);
    printf("    i: %.3f %.3f %.3f\n",
           incoming_color.x_,
           incoming_color.y_,
           incoming_color.z_);
    printf("    il: %.3f %.3f %.3f\n",
           incoming_light_color.x_,
           incoming_light_color.y_,
           incoming_light_color.z_);
//    assert(false);
  }
  */

  return path_color;
}

Scene::PrimitiveList::const_iterator Scene::Intersect(const Ray& r,
                                                     float& intersect_d) const {
  PrimitiveList::const_iterator min_prim_it = primitives_.end();
  intersect_d = std::numeric_limits<float>::infinity();

  for (PrimitiveList::const_iterator it = primitives_.begin();
       it != primitives_.end(); ++it) {
    float d = 0;
    if ((*it)->IntersectsWith(r, d)) {
      if (d < intersect_d && d > 0) {
        intersect_d = d;
        min_prim_it = it;
      }
    }
  }

  return min_prim_it;
}


} // end namespace
